#![allow(non_snake_case)]

use c2rust003::{chatgml::{sse_invoke_method, C2RustGLM}, code2prompt::codes2prompts, prompt::generate_prompt};
use c2rust004::{CodeFile, TranspileFile};

// use serde::{Deserialize, Serialize};
use std::fs;
use std::fs::File;
mod category_parse;

use regex::Regex;
use sse_invoke_method::sse_invoke::constant_value::API_KEY;
use std::collections::HashMap;
#[allow(unused_imports)]
use std::io::{self, Write};
#[allow(unused_imports)]
use std::path::Path;

pub fn remove_markdown_code_block(text: &str) -> String {
    // 创建正则表达式来匹配代码块标记
    let re = Regex::new(r"```.*?\n").unwrap();
    let text = re.replace_all(text, "");

    // 去除剩余的代码块结束标记
    let re = Regex::new(r"```.*?").unwrap();
    let text = re.replace_all(&text, "");

    text.to_string()
}

// TODO: 这里需要修改，暂时这样写
pub fn traverse_dir(
    dir_path: &str,
    test_path: &str,
    filenames: Vec<String>,
    header_files: &mut HashMap<String, String>,
    source_files: &mut HashMap<String, String>,
    test_files: &mut HashMap<String, String>,
) -> std::io::Result<()> {
    let dir = fs::read_dir(dir_path)?;
    for entry in dir {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            //println!("File name:{:?}",path.file_name().unwrap());
            if let Some(extension) = path.extension() {
                if extension == "h" {
                    // 读取头文件内容
                    if let Ok(content) = fs::read_to_string(&path) {
                        let filename = path.file_stem().unwrap().to_str().unwrap().to_string();
                        if filenames.contains(&filename) {
                            header_files.insert(filename.clone(), content);
                        }
                    }
                } else if extension == "c" {
                    // 读取C源文件内容
                    if let Ok(content) = fs::read_to_string(&path) {
                        let filename = path.file_stem().unwrap().to_str().unwrap().to_string();
                        if filenames.contains(&filename) {
                            source_files.insert(filename.clone(), content);
                        }
                    }
                }
            }
        } else if path.is_dir() {
            let _ = traverse_dir(
                &path.to_string_lossy(),
                test_path,
                filenames.clone(),
                header_files,
                source_files,
                test_files,
            )
            .unwrap();
        }
    }

    let test_dir = fs::read_dir(test_path)?;
    let test_names = filenames
        .iter()
        .map(|it| format!("test-{}", it))
        .collect::<Vec<String>>();
    for entry in test_dir {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            //println!("File name:{:?}",path.file_name().unwrap());
            if let Ok(content) = fs::read_to_string(&path) {
                let filename = path.file_stem().unwrap().to_str().unwrap().to_string();
                if test_names.contains(&filename) {
                    test_files.insert(filename.clone(), content);
                }
            }
        } else if path.is_dir() {
            let _ = traverse_dir(
                &path.to_string_lossy(),
                test_path,
                filenames.clone(),
                header_files,
                source_files,
                test_files,
            )
            .unwrap();
        }
    }

    Ok(())
}

// pub struct CodeFile {
//     pub filename: String,
//     pub content: String,
//     pub test: Option<String>,
// }

pub fn translate_prompt(base_dir: &str, filenames: Vec<String>) -> Vec<CodeFile> {
    let mut ret = Vec::<CodeFile>::new();
    // 指定目录
    let binding = ["./Input/", &base_dir, "/src"].concat();
    let src_directory: &str = binding.as_str();
    let binding1 = ["./Input/", &base_dir, "/test"].concat();
    let test_directory: &str = binding1.as_str();

    let base_dir = ["./Input/", &base_dir].concat();

    let _rendered = codes2prompts(&base_dir).unwrap();
    // println!("函数接受回来的数据： rendered: {}", rendered);

    // 读取目录中的所有文件
    let mut header_files = HashMap::new();
    let mut source_files = HashMap::new();
    let mut test_files = HashMap::new();

    let _ = traverse_dir(
        src_directory,
        test_directory,
        filenames,
        &mut header_files,
        &mut source_files,
        &mut test_files,
    )
    .unwrap();

    // 处理文件对
    for (filename, header_content) in &header_files {
        if let Some(source_content) = source_files.get(filename) {
            // 生成转译代码的prompt
            let mut test_prompt = None;
            if let Some(c) = test_files.get(format!("test-{}", filename).as_str()) {
                let rust_filename = filename.clone().replace(".c", "rs");
                let prompt = format!("Translate the following C code to Rust, Output only the code content, do not need to output the comment content; code is {}; This should be a main.rs file content and import the public function in this file,it's filename is : {}, the generated code on last step and the current source code is: {}, ", c, rust_filename, source_content);
                let prompt = generate_prompt(&prompt);
                test_prompt = Some(prompt);
            }
            let prompt = format!("Translate the following C code to Rust; Output only the code content, do not need to output the comment content; and do not need output the main function in the code but need to public the function; this is head file content: {} \n and this is source file content: {}, ", header_content, source_content);
            let prompt = generate_prompt(&prompt);
            ret.push(CodeFile {
                filename: filename.clone(),
                content: prompt,
                test: test_prompt,
            });
        }
    }
    ret
}

pub async fn zhipu_transpile(user_input: &str) -> String {
    let mut c_2_rust_glm = C2RustGLM::new().await;
    c_2_rust_glm.sse_invoke_calling(API_KEY, user_input).await;
    let result = c_2_rust_glm.get_ai_response();
    result
}

pub async fn transpile_file_s(file: TranspileFile) -> TranspileFile {
    // TODO: 实现文件转译逻辑，可以适配多种转译模式
    // 调用API进行转译
    let content = zhipu_transpile(&file.content).await;
    TranspileFile {
        filename: file.filename,
        content,
    }
}

// 生成Rust项目
pub async fn generate_rust_project(name: String, c_2_rust_glm: &mut C2RustGLM, prompt: CodeFile) {
    let src_dir = format!("./Output/{}", name.clone());
    let directory: &str = &src_dir.as_str();
    let src_d = format!("./Output/{}/src", name.clone());
    let src_str = &src_d.as_str();
    // 判断目录是否存在，没有话创建
    if !Path::new(directory).exists() {
        fs::create_dir_all(directory).expect("Failed to create directory");
        fs::create_dir_all(src_str).expect("Failed to create directory");
    }

    let mut toml =
        File::create(format!("{}/{}", directory, "Cargo.toml")).expect("Failed to create file");
    let toml_str = format!(
        r#"
[package]
name="{}"
version = "0.0.1"
edition = "2021"
[dependencies]

[[test]]
name = "tests"
path = "src/tests.rs"
    "#,
        name
    );
    toml.write_all(toml_str.as_bytes())
        .expect("Failed to write to file");
    let filename = name.clone() + ".rs"; // "main.rs";
    let mut file =
        File::create(format!("{}/src/{}", directory, filename)).expect("Failed to create file");
    let mut test =
        File::create(format!("{}/src/{}", directory, "main.rs")).expect("Failed to create file");

    let result = c_2_rust_glm.get_ai_response();
    println!("{}", result);
    file.write_all(remove_markdown_code_block(&result).as_bytes())
        .expect("Failed to write to file");
    if let Some(t) = &prompt.test {
        c_2_rust_glm.sse_invoke_calling(API_KEY, t).await;
        let test_result = c_2_rust_glm.get_ai_response();
        test.write_all(remove_markdown_code_block(&test_result).as_bytes())
            .expect("Failed to write to file");
    }
}
